package com.mega.fe.mixins;

import com.google.common.collect.ImmutableList;
import com.google.common.collect.Iterators;
import com.mega.fe.util.NonNullIterable;
import net.minecraft.util.ClassInstanceMultiMap;
import org.spongepowered.asm.mixin.Final;
import org.spongepowered.asm.mixin.Mixin;
import org.spongepowered.asm.mixin.Shadow;
import org.spongepowered.asm.mixin.injection.At;
import org.spongepowered.asm.mixin.injection.Inject;
import org.spongepowered.asm.mixin.injection.callback.CallbackInfo;
import org.spongepowered.asm.mixin.injection.callback.CallbackInfoReturnable;

import java.util.*;
import java.util.stream.Collectors;

@Mixin(ClassInstanceMultiMap.class)
public class ClassInstanceMultiMapCreateListener<T> {
    @Shadow @Final private Class<T> baseClass;

    @Shadow @Final private Map<Class<?>, List<T>> byClass;

    @Shadow @Final private List<T> allInstances;

    @Inject(method = "<init>", at = @At("RETURN"))
    private void init(Class p_13531_, CallbackInfo ci) {
    }


    @Inject(method = "find", at = @At("HEAD"), cancellable = true)
    private<S> void find(Class<S> p_13534_, CallbackInfoReturnable<Collection<S>> cir) {
        if (!baseClass.isAssignableFrom(p_13534_)) {
            throw new IllegalArgumentException("Don't know how to search for " + p_13534_);
        } else {
            List<? extends T> list = byClass.computeIfAbsent(p_13534_, (p_13538_) -> {
                List<T> l =  allInstances.stream()
                        .filter(p_13538_::isInstance)
                        .collect(Collectors.toList());
                l.removeAll(Collections.singleton(null));
                return l;
            });
            cir.setReturnValue( (Collection<S>) Collections.unmodifiableCollection(list));
        }
    }
    @Inject(method = "iterator", at = @At("HEAD"), cancellable = true)
    private void iterator(CallbackInfoReturnable<Iterator<T>> cir) {
        cir.setReturnValue((allInstances.isEmpty() ? Collections.emptyIterator() : Iterators.unmodifiableIterator(new NonNullIterable<>(allInstances).iterator())));
    }

    @Inject(method = "getAllInstances", at = @At("HEAD"), cancellable = true)
    private void getAllInstances(CallbackInfoReturnable<List<T>> cir) {
        List<T> list = allInstances;
        list.removeAll(Collections.singleton(null));
        cir.setReturnValue( ImmutableList.copyOf(list));
    }
}
